import gradio as gr
import torch
import torch.nn.functional as F
from transformers import BertTokenizer, GPT2LMHeadModel

import time
import multiprocessing
import subprocess
import psutil

print(f"Is CUDA available: {torch.cuda.is_available()}")
# True
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
# Tesla T4



tokenizer = BertTokenizer.from_pretrained("hf-models/gpt2-chinese-poem")
model = GPT2LMHeadModel.from_pretrained("hf-models/gpt2-chinese-poem")
model.eval()


def print_gpu():
    # 打印 GPU 信息
    result = subprocess.run(['ixsmi'], capture_output=True, text=True)
    print('ixsmi Output:', result.stdout)

    # 打印CPU信息
    print("CPU信息:")
    cpu_count = psutil.cpu_count(logical=False)  # 获取物理CPU数量
    print(f"物理CPU数量: {cpu_count}")
    cpu_count_logical = psutil.cpu_count()  # 获取逻辑CPU数量
    print(f"逻辑CPU数量: {cpu_count_logical}")
    cpu_freq = psutil.cpu_freq()  # 获取CPU频率
    print(f"CPU频率: {cpu_freq.current} MHz")

    # 打印内存信息
    print("\n内存信息:")
    memory = psutil.virtual_memory()
    print(f"总内存: {memory.total / (1024 ** 3):.2f} GB")
    print(f"可用内存: {memory.available / (1024 ** 3):.2f} GB")
    print(f"已使用内存: {memory.used / (1024 ** 3):.2f} GB")
    print(f"内存使用率: {memory.percent}%")

    # 打印当前CPU使用率
    print("\n当前CPU使用率:")
    cpu_usage = psutil.cpu_percent(interval=1)  # interval参数表示间隔时间，单位为秒
    print(f"CPU使用率: {cpu_usage}%")

    # 打印当前内存使用率
    print("\n当前内存使用率:")
    memory_usage = memory.percent
    print(f"内存使用率: {memory_usage}%")


def print_gpu_process(event, stop_event):
    while not stop_event.is_set():  # 检查事件是否被设置
        if event.is_set():
            print_gpu()
        time.sleep(2)
    print("子进程结束")


def start_print_gpu_process(event, stop_event):
    # 创建一个新的进程
    process = multiprocessing.Process(target=print_gpu_process, args=(event, stop_event,))
    # 启动子进程
    process.start()
    print("启动子进程")
    return process


print_event = multiprocessing.Event()
# 创建一个事件，用于通知子进程何时停止
stop_event = multiprocessing.Event()



def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
    assert logits.dim() == 1
    top_k = min(top_k, logits.size(-1))
    if top_k > 0:
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value
    if top_p > 0.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[indices_to_remove] = filter_value
    return logits





def generate(input_text):
    input_ids = [tokenizer.cls_token_id]
    input_ids.extend( tokenizer.encode(input_text, add_special_tokens=False) )
    input_ids = torch.tensor( [input_ids] )
    generated = []
    # 启动打印 gpu 信息
    print_event.set()
    for _ in range(100):
        output = model(input_ids)
        next_token_logits = output.logits[0, -1, :]
        next_token_logits[tokenizer.convert_tokens_to_ids('[UNK]')] = -float('Inf')
        filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=8, top_p=1)
        next_token = torch.multinomial( F.softmax(filtered_logits, dim=-1), num_samples=1 )
        if next_token == tokenizer.sep_token_id:
            break
        generated.append( next_token.item() )
        input_ids = torch.cat((input_ids, next_token.unsqueeze(0)), dim=1)
    # 停止打印gpu信息
    print_event.clear()
    return input_text + "".join( tokenizer.convert_ids_to_tokens(generated) )








examples = [["不堪翘首暮云中"], ["开源中国"], ["行到水穷处"], ["王师北定中原日"] ,["雪"], ["海上升明月"], ["十年磨一剑"]]







if __name__ == "__main__":

    # 启动子进程
    print_process = start_print_gpu_process(print_event, stop_event)

    # gr.Interface(fn=generate, inputs="text", outputs="text",examples=examples).queue(concurrency_count=1).launch()
    demo = gr.Interface(fn=generate, inputs="text", outputs="text", examples=examples)
    demo.launch(server_name='0.0.0.0')